Permit `BudgetOptimizer.allocate_budget()` to take `x0` as an argument #1565

cluhmann · 2025-03-16T21:22:30Z

Description

BudgetOptimizer.allocate_budget() now accepts x0 as an argument.

Related Issue

Closes BudgetOptimizer.allocate_budget() doc string doesn't specify default value of x0 #1564
Related to #

Checklist

Checked that the pre-commit linting/style checks pass. Feel free to comment pre-commit.ci autofix to auto-fix.
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks) using numpydoc format.
If you are a pro: each commit corresponds to a relevant logical change

x0 is now accepted via the minimize_kwargs (as the docstring implies is possible). I looked into adding a corresponding test, but these lines in tests/mmm/test_budget_optimizer::test_allocate_budget_custom_minimize_args() confuse me. Any guidance?

kwargs = minimize_mock.call_args_list[0].kwargs

np.testing.assert_array_equal(x=kwargs["x0"], y=np.array([50.0, 50.0]))
assert kwargs["bounds"] == [(0.0, 50.0), (0.0, 50.0)]
assert kwargs["method"] == minimize_kwargs["method"]
assert kwargs["options"] == minimize_kwargs["options"]

📚 Documentation preview 📚: https://pymc-marketing--1565.org.readthedocs.build/en/1565/

codecov · 2025-03-16T21:25:20Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 93.53%. Comparing base (a87f1fa) to head (63c24ed).
Report is 4 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1565   +/-   ##
=======================================
  Coverage   93.53%   93.53%           
=======================================
  Files          55       55           
  Lines        6357     6358    +1     
=======================================
+ Hits         5946     5947    +1     
  Misses        411      411

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2025-03-17T17:10:13Z

pymc_marketing/mmm/budget_optimizer.py

-            # Merge with defaults (preferring user-supplied keys)
-            minimize_kwargs = {**self.DEFAULT_MINIMIZE_KWARGS, **minimize_kwargs}
+        # 5. Construct the initial guess (x0) if not provided
+        if "x0" not in minimize_kwargs:


I don't think x0 should be hidden in minimize kwargs, I would make it an explicit argument of allocate_budget

ricardoV94 · 2025-03-17T17:11:42Z

pymc_marketing/mmm/budget_optimizer.py

+            minimize_kwargs["x0"] = np.ones(budgets_size) * (
+                total_budget / budgets_size
+            )
+            minimize_kwargs["x0"] = minimize_kwargs["x0"].astype(


I would check if it has the expected shape (budgets_size,), as the compiled functions have trust_input=True and won't check if x0 has the promised shape

very good point here!

cetagostini · 2025-03-17T21:09:51Z

kwargs = minimize_mock.call_args_list[0].kwargs
np.testing.assert_array_equal(x=kwargs["x0"], y=np.array([50.0, 50.0]))
assert kwargs["bounds"] == [(0.0, 50.0), (0.0, 50.0)]
assert kwargs["method"] == minimize_kwargs["method"]
assert kwargs["options"] == minimize_kwargs["options"]

Just a hardcoded check based on the data, we could change this to be more dynamic and check if values are equal based on the data, not hardcoded 50/50 🙃 Does that makes it clear?

ricardoV94 · 2025-03-19T11:00:17Z

pymc_marketing/mmm/budget_optimizer.py

        self._total_budget.set_value(np.asarray(total_budget, dtype="float64"))

+        # coordinate user-provided and default minimize_kwargs
+        if minimize_kwargs is None:
+            minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS.copy()


Nit: no real need to copy

Suggested change

minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS.copy()

minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS

ricardoV94 · 2025-03-19T11:00:42Z

pymc_marketing/mmm/budget_optimizer.py

@@ -391,8 +392,11 @@ def allocate_budget(
            - If None, default bounds of [0, total_budget] per channel are assumed.
            - If a dict, must map each channel to (low, high) budget pairs (only valid if there's one dimension).
            - If an xarray.DataArray, must have dims (*budget_dims, "bound"), specifying [low, high] per channel cell.
+        x0 : DataArray, optional


I think it should be a np.ndarray not a DataArray

ricardoV94 · 2025-03-19T11:01:14Z

pymc_marketing/mmm/budget_optimizer.py

+                self._budgets_flat.type.dtype
+            )
+        # if x0 arg is provided, validate shape
+        elif x0.shape != (budgets_size,):


Should also validate dtype

cluhmann · 2025-03-19T12:23:14Z

Just need to write some tests.

ricardoV94 · 2025-03-24T10:07:45Z

pymc_marketing/mmm/budget_optimizer.py

+                f"""The shape of 'x0' {x0.shape} does not match the expected shape {(budgets_size,)}."""
+            )
+        # if x0 arg is provided, validate dtype
+        elif np.issubdtype(x0.dtype, np.number):


Specifically the x0.dtype must match self._budgets_flat.type.dtype. I think there's a shortcut you can use for the shape and dtype validation: x0 = self._budgets_flat.type.filter(x0), which will do any allowed casting or raise if things cannot be safely cast (and also checks shape)

cetagostini

Everything looks good from my side, all mentions from Ricardo are address and the test are working. After a look to the code everything is ready to merge, if something is missing we can get it done in another PR.

cluhmann · 2025-04-15T20:50:38Z

Thank you! If no one else wants to merge this, I will at some point soon.

williambdean · 2025-04-15T21:17:40Z

Thank you! If no one else wants to merge this, I will at some point soon.

do the honors!

cluhmann requested review from ricardoV94 and cetagostini March 16, 2025 21:22

github-actions bot added MMM optimizer labels Mar 16, 2025

ricardoV94 reviewed Mar 17, 2025

View reviewed changes

juanitorduz requested a review from ricardoV94 March 18, 2025 20:20

ricardoV94 reviewed Mar 19, 2025

View reviewed changes

ricardoV94 reviewed Mar 24, 2025

View reviewed changes

github-actions bot added the tests label Mar 25, 2025

cluhmann added 5 commits March 25, 2025 16:48

formatting

7b6fd91

expose x0 as an arg in allocate_budget

0427f1f

change typehint and validate dtype

85d6ccb

streamline shape and dtype validation

d3919bf

add new bo test with manually set x0

6683395

cluhmann force-pushed the bo_x0 branch from a94c35a to 6683395 Compare March 25, 2025 20:50

cetagostini added 2 commits March 26, 2025 13:33

Merge branch 'main' into bo_x0

c8a28ec

Merge branch 'main' into bo_x0

63c24ed

cetagostini approved these changes Apr 15, 2025

View reviewed changes

cluhmann merged commit a891405 into main Apr 15, 2025
33 checks passed

williambdean deleted the bo_x0 branch April 15, 2025 21:39

cluhmann mentioned this pull request Apr 21, 2025

Fix casting bug when x0 is not supplied to optimizer #1636

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Permit `BudgetOptimizer.allocate_budget()` to take `x0` as an argument #1565

Permit `BudgetOptimizer.allocate_budget()` to take `x0` as an argument #1565

Uh oh!

cluhmann commented Mar 16, 2025 •

edited

Loading

Uh oh!

codecov bot commented Mar 16, 2025 •

edited

Loading

Uh oh!

ricardoV94 Mar 17, 2025

Uh oh!

ricardoV94 Mar 17, 2025

Uh oh!

cetagostini Mar 17, 2025

Uh oh!

cetagostini commented Mar 17, 2025

Uh oh!

ricardoV94 Mar 19, 2025

Uh oh!

ricardoV94 Mar 19, 2025

Uh oh!

ricardoV94 Mar 19, 2025

Uh oh!

cluhmann commented Mar 19, 2025

Uh oh!

ricardoV94 Mar 24, 2025 •

edited

Loading

Uh oh!

cetagostini left a comment

Uh oh!

cluhmann commented Apr 15, 2025

Uh oh!

williambdean commented Apr 15, 2025

Uh oh!

Uh oh!

Uh oh!

	minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS.copy()
	minimize_kwargs = self.DEFAULT_MINIMIZE_KWARGS

Permit BudgetOptimizer.allocate_budget() to take x0 as an argument #1565

Permit BudgetOptimizer.allocate_budget() to take x0 as an argument #1565

Uh oh!

Conversation

cluhmann commented Mar 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Uh oh!

codecov bot commented Mar 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 Mar 17, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 17, 2025

Choose a reason for hiding this comment

Uh oh!

cetagostini Mar 17, 2025

Choose a reason for hiding this comment

Uh oh!

cetagostini commented Mar 17, 2025

Uh oh!

ricardoV94 Mar 19, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 19, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 19, 2025

Choose a reason for hiding this comment

Uh oh!

cluhmann commented Mar 19, 2025

Uh oh!

ricardoV94 Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cetagostini left a comment

Choose a reason for hiding this comment

Uh oh!

cluhmann commented Apr 15, 2025

Uh oh!

williambdean commented Apr 15, 2025

Uh oh!

Uh oh!

Uh oh!

Permit `BudgetOptimizer.allocate_budget()` to take `x0` as an argument #1565

Permit `BudgetOptimizer.allocate_budget()` to take `x0` as an argument #1565

cluhmann commented Mar 16, 2025 •

edited

Loading

codecov bot commented Mar 16, 2025 •

edited

Loading

ricardoV94 Mar 24, 2025 •

edited

Loading